// Copyright (c) 2020 cgsdfc
//
// This software is released under the MIT License.
// https://opensource.org/licenses/MIT

#pragma once

#include <memory>

#include "pyast/pybind_defs.h"

namespace pyast {

// Shell class around PyArena struct.
class Arena {
 public:
  struct DeleteArenaFunctor {
    void operator()(Arena* p) const noexcept {  // Calling C code won't throw.
      if (p)
        Arena::DeleteArenaImpl(p);  //
    }
  };

  using SharedPtr = std::shared_ptr<Arena>;
  using UniquePtr = std::unique_ptr<Arena, DeleteArenaFunctor>;

  // Create an Arena with shared ownership. The lifetime of the created Arena is
  // managed by `std::shared_ptr` with a custom `Delete` functor.
  static SharedPtr CreateShared() {
    return SharedPtr{CreateRaw(), DeleteArenaFunctor{}};  //
  }

  // Create an Arena with unique ownership. The lifetime of the created Arena is
  // managed by `std::unique_ptr` with a custom `Delete` functor.
  static UniquePtr CreateUnique() {
    return UniquePtr{CreateRaw()};  //
  }

  static Arena* CreateRaw() {
    Arena* new_arena = CreateArenaImpl();
    if (!new_arena)
      throw error_already_set{};
    return new_arena;
  }

  void* Malloc(size_t size) {
    if (void* result = ArenaMallocImpl(this, size))
      return result;
    throw error_already_set();
  }

  void AddObject(object obj) {
    if (!obj)
      return;
    if (ArenaAddObjectImpl(this, std::move(obj)))
      throw error_already_set{};
  }

  // Self-deleting routine.
  // Occationally you can obtain a raw pointer to Arena, when you can use this
  // method to manually delete the Arena object. But mostly you rely on the
  // smart pointers to do that for you. Well, mostly.
  void Delete() noexcept {
    DeleteArenaImpl(this);  //
  }

  PyArena* ptr() { return ToPyArena(this); }

 private:
  // Destructor should never get a chance to be called.
  ~Arena() = delete;
  // Constructor should never get a chance to be called.
  Arena() = delete;

  // Convert self to the C's counterpart.
  static PyArena* ToPyArena(Arena* arena) {
    return reinterpret_cast<PyArena*>(arena);  //
  }

  // Implement other things in terms of these routines.
  // These are simple C code wrappers so that the C code does not fly
  // everywhere.
  static void DeleteArenaImpl(Arena* arena) noexcept {
    return PyArena_Free(ToPyArena(arena));
  }

  static Arena* CreateArenaImpl() noexcept {
    return reinterpret_cast<Arena*>(PyArena_New());
  }

  static void* ArenaMallocImpl(Arena* arena, size_t size) noexcept {
    return PyArena_Malloc(ToPyArena(arena), size);
  }

  static int ArenaAddObjectImpl(Arena* arena, object&& obj) noexcept {
    // The arena takes ownership of `obj`.
    return PyArena_AddPyObject(ToPyArena(arena), obj.release().ptr());
  }
};

// To avoid typing `std::shared_ptr<Arena>` or `std::unique_ptr<Arena>`
// repeatedly (which is tedious), use these aliases for better life.
// If you need a raw pointer, just use `Arena*`.
using SharedArena = Arena::SharedPtr;
using UniqueArena = Arena::UniquePtr;

}  // namespace pyast
